#include <stdio.h>
#include <NvInfer.h>
#include <NvOnnxParser.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cuda_runtime_api.h>
#include <iostream>
#include <fstream>
#include <opencv2/opencv.hpp>
#include <opencv2/dnn.hpp>
#include "trthelper.h"
#include "yolov5_deploy.h"
#include "trtModelDeploy.h"
#include "factory.h"
using namespace nvonnxparser;
using namespace nvinfer1;




std::vector<std::string> labels = TRTHelper::getInstance().readClassNames("/home/xtdz/project/model/classes.txt");



int main(int argc, char** argv)
{
    Logger logger;
    printf("hello, world\n");
    if(argc < 2)
    {
        std::cout << "usage: ./exec model_path" << std::endl;
        return 0;
    }

    

    auto &config_manager = ModelConfigManager::getInstance();   
    config_manager.loadConfig("/home/xtdz/project/trt-model-deploy/bin/config.txt");
    const auto& config = config_manager.getConfig(argv[1]);
    
    std::string model_name = argv[1];
    std::unique_ptr<trtModelDeploy> trtModelInfer;
    std::cout << "the config is " << config.image_path << std::endl;
    std::cout << "the model_path is " << config.model_path << std::endl;
    cv::Mat image = cv::imread(config.image_path);
    if(model_name == "yolov5") 
    {
        trtModelInfer = ModelFactory::createModel(YOLOV5, config.model_path, config.classes_path);
        
    }else if(model_name == "yolov8")
    {
        trtModelInfer = ModelFactory::createModel(YOLOV8, config.model_path, config.classes_path);
    }else if(model_name == "yolov11")
    {
        trtModelInfer = ModelFactory::createModel(YOLOV11, config.model_path, config.classes_path);
    }else if(model_name == "yolov5-seg")
    {
        trtModelInfer = ModelFactory::createModel(YOLOV5_SEG, config.model_path, config.classes_path);
    }else if(model_name == "yolov8-seg")
    {
        trtModelInfer = ModelFactory::createModel(YOLOV8_SEG, config.model_path, config.classes_path);
    }else if(model_name == "yolov11-seg")
    {
        trtModelInfer = ModelFactory::createModel(YOLOV11_SEG, config.model_path, config.classes_path);
    }else if (model_name == "yolov8-pose")
    {
        trtModelInfer = ModelFactory::createModel(YOLOV8_POSE, config.model_path, config.classes_path);
    }else if (model_name == "yolov11-pose")
    {
        trtModelInfer = ModelFactory::createModel(YOLOV11_POSE, config.model_path, config.classes_path);
    }else if (model_name == "yolov6-face")
    {
        trtModelInfer = ModelFactory::createModel(YOLOV6_FACE, config.model_path, config.classes_path);
    }
    else
    {
        std::cout << "model not found" << std::endl;
    }

    if(trtModelInfer)
    {
        cv::Mat pre_blob = trtModelInfer->preprocess(image);
        trtModelInfer->infer(pre_blob);
        trtModelInfer->postprocess(image);
    }
    
    

    return 0;
}